#!/usr/bin/env python3

"""
Lerobot Bridge Node - 桥接 ROS2 Control 和 lerobot 库
"""

import rclpy
from rclpy.node import Node
from sensor_msgs.msg import JointState
from std_srvs.srv import SetBool
import numpy as np
import threading
import sys
import os

# 添加 lerobot 路径（根据你的安装位置调整）
lerobot_path = "/home/parallels/PycharmProjects/aghanim-scepter/references/lerobot"
if lerobot_path not in sys.path:
    sys.path.append(lerobot_path)

try:
    from robot_interface import ManipulatorRobot
    from lerobot.common.robot_devices.robots.configs import So101RobotConfig
except ImportError as e:
    print(f"Failed to import lerobot modules: {e}")
    print("Please check your lerobot installation path")
    sys.exit(1)


class LerobotBridgeNode(Node):
    """桥接 ROS2 Control 和 lerobot 库的节点"""
    
    def __init__(self):
        super().__init__('lerobot_bridge_node')
        
        # 声明参数
        self.declare_parameter('robot_config', 'so101')
        self.declare_parameter('update_rate', 100.0)  # Hz
        self.declare_parameter('command_timeout', 1.0)  # seconds
        
        # 获取参数
        robot_config_name = self.get_parameter('robot_config').value
        self.update_rate = self.get_parameter('update_rate').value
        self.command_timeout = self.get_parameter('command_timeout').value
        
        self.get_logger().info(f"Initializing with robot config: {robot_config_name}")
        
        # 初始化机器人
        try:
            if robot_config_name == 'so101':
                config = So101RobotConfig()
            else:
                self.get_logger().error(f"Unsupported robot config: {robot_config_name}")
                return
                
            self.robot = ManipulatorRobot(config)
            self.get_logger().info("Robot interface created successfully")
        except Exception as e:
            self.get_logger().error(f"Failed to create robot interface: {e}")
            return
        
        # 状态变量
        self.robot_connected = False
        self.robot_enabled = False
        self.last_command_time = self.get_clock().now()
        self.command_lock = threading.Lock()
        
        # 关节信息
        self.joint_names = ['joint1', 'joint2', 'joint3', 'joint4', 'joint5', 'joint6', 'hand_right_joint']
        self.num_joints = len(self.joint_names)
        self.current_positions = np.zeros(self.num_joints)
        self.current_velocities = np.zeros(self.num_joints)
        self.target_positions = np.zeros(self.num_joints)
        
        # ROS 接口
        self.setup_ros_interfaces()
        
        # 定时器
        timer_period = 1.0 / self.update_rate  # Convert Hz to seconds
        self.timer = self.create_timer(timer_period, self.control_loop)
        
        self.get_logger().info(f"Lerobot bridge node initialized (update rate: {self.update_rate} Hz)")
    
    def setup_ros_interfaces(self):
        """设置 ROS 接口"""
        
        # 发布关节状态给硬件接口
        self.joint_state_publisher = self.create_publisher(
            JointState, '/lerobot/joint_states', 10)
        
        # 订阅来自硬件接口的命令
        self.command_subscriber = self.create_subscription(
            JointState, '/lerobot/joint_commands', self.command_callback, 10)
        
        # 服务：启用/禁用机器人
        self.enable_service = self.create_service(
            SetBool, '/lerobot/enable_robot', self.enable_robot_callback)
        
        # 服务：连接/断开机器人
        self.connect_service = self.create_service(
            SetBool, '/lerobot/connect_robot', self.connect_robot_callback)
        
        self.get_logger().info("ROS interfaces created")
    
    def connect_robot_callback(self, request, response):
        """连接/断开机器人服务回调"""
        try:
            if request.data:
                if not self.robot_connected:
                    self.robot.connect()
                    self.robot_connected = True
                    # 读取当前位置
                    current_pos = self.robot.get_positions(degree=False)
                    self.current_positions[:len(current_pos)] = current_pos
                    self.target_positions = self.current_positions.copy()
                    
                    self.get_logger().info("Robot connected successfully")
                else:
                    self.get_logger().info("Robot already connected")
                response.success = True
                response.message = "Robot connected"
            else:
                if self.robot_connected:
                    # 断开连接时先禁用电机
                    if self.robot_enabled:
                        self.robot.disable_motor()
                        self.robot_enabled = False
                    # 注意：lerobot 可能没有显式的断开方法
                    self.robot_connected = False
                    self.get_logger().info("Robot disconnected")
                response.success = True
                response.message = "Robot disconnected"
                
        except Exception as e:
            self.get_logger().error(f"Failed to connect/disconnect robot: {e}")
            response.success = False
            response.message = str(e)
            
        return response
    
    def enable_robot_callback(self, request, response):
        """启用/禁用机器人服务回调"""
        try:
            if not self.robot_connected:
                response.success = False
                response.message = "Robot not connected"
                return response
                
            if request.data:
                if not self.robot_enabled:
                    self.robot.enable_motor()
                    self.robot_enabled = True
                    self.get_logger().info("Robot motors enabled")
                else:
                    self.get_logger().info("Robot motors already enabled")
                response.success = True
                response.message = "Robot enabled"
            else:
                if self.robot_enabled:
                    self.robot.disable_motor()
                    self.robot_enabled = False
                    self.get_logger().info("Robot motors disabled")
                response.success = True
                response.message = "Robot disabled"
                
        except Exception as e:
            self.get_logger().error(f"Failed to enable/disable robot: {e}")
            response.success = False
            response.message = str(e)
            
        return response
    
    def command_callback(self, msg):
        """接收来自硬件接口的命令"""
        if not self.robot_connected or not self.robot_enabled:
            return
            
        with self.command_lock:
            # 更新目标位置
            for i, joint_name in enumerate(msg.name):
                if joint_name in self.joint_names and i < len(msg.position):
                    joint_idx = self.joint_names.index(joint_name)
                    self.target_positions[joint_idx] = msg.position[i]
            
            self.last_command_time = self.get_clock().now()
    
    def control_loop(self):
        """主控制循环"""
        if not self.robot_connected:
            return
            
        try:
            # 读取当前关节状态
            current_pos = self.robot.get_positions(degree=False)
            if len(current_pos) <= len(self.current_positions):
                self.current_positions[:len(current_pos)] = current_pos
            
            # 估计速度（简单差分）
            # TODO: 可以实现更好的速度估计
            # self.current_velocities = ... 
            
            # 发布关节状态
            self.publish_joint_states()
            
            # 发送命令到机器人（如果有新命令且在超时范围内）
            if self.robot_enabled:
                current_time = self.get_clock().now()
                time_diff = (current_time - self.last_command_time).nanoseconds * 1e-9
                
                if time_diff < self.command_timeout:
                    with self.command_lock:
                        # 使用安全的目标位置
                        safe_positions = self.ensure_safe_positions(self.target_positions)
                        self.robot.set_positions(safe_positions, degree=False)
                
        except Exception as e:
            self.get_logger().error(f"Control loop error: {e}")
    
    def ensure_safe_positions(self, target_positions):
        """确保目标位置安全"""
        # 实现位置限制和安全检查
        safe_positions = target_positions.copy()
        
        # 这里可以添加关节限制检查
        # 例如：safe_positions = np.clip(safe_positions, min_limits, max_limits)
        
        return safe_positions
    
    def publish_joint_states(self):
        """发布关节状态"""
        msg = JointState()
        msg.header.stamp = self.get_clock().now().to_msg()
        msg.name = self.joint_names
        msg.position = self.current_positions.tolist()
        msg.velocity = self.current_velocities.tolist()
        
        self.joint_state_publisher.publish(msg)


def main(args=None):
    rclpy.init(args=args)
    
    try:
        node = LerobotBridgeNode()
        
        # 自动连接机器人
        try:
            if hasattr(node, 'robot'):
                node.robot.connect()
                node.robot_connected = True
                # 读取初始位置
                current_pos = node.robot.get_positions(degree=False)
                node.current_positions[:len(current_pos)] = current_pos
                node.target_positions = node.current_positions.copy()
                node.get_logger().info("Robot auto-connected successfully")
        except Exception as e:
            node.get_logger().warning(f"Auto-connect failed: {e}")
        
        rclpy.spin(node)
        
    except KeyboardInterrupt:
        pass
    except Exception as e:
        print(f"Error: {e}")
    finally:
        # 清理
        try:
            if 'node' in locals() and hasattr(node, 'robot') and node.robot_connected:
                if node.robot_enabled:
                    node.robot.disable_motor()
                node.get_logger().info("Robot safely disconnected")
        except Exception as e:
            print(f"Cleanup error: {e}")
        
        rclpy.shutdown()


if __name__ == '__main__':
    main()